import core
import numpy
import math


#pyfitter=core.pyfitter
#pychisq=core.pychisq
from core import pyfitter
from core import pychisq
from core import pybootstrap



def create_fitter():
    result=pyfitter()
    result.set_statistic(core.pychisq())
    result.set_opt_method(core.pypowell_method())
    return result



def create_data_set(x,y,ye1,ye2,xe1,xe2):
    data_size=len(x)
    result=core.pysorted_data_set()
    for i in range(0,data_size):
        result.add_data(core.pydata(x[i],y[i],ye1[i],ye2[i],xe1[i],xe2[i]))
    
    return result


def load_data(f):
    result=core.pysorted_data_set();
    
    for i in f:
        x,xe1,xe2,y,ye1,ye2=i.split()
        x=float(x)
        y=float(y)
        xe1=math.fabs(float(xe1))
        xe2=math.fabs(float(xe2))
        ye1=math.fabs(float(ye1))
        ye2=math.fabs(float(ye2))

        result.add_data(core.pydata(x,y,ye1,ye2,xe1,xe2))
    
    return result;

def load_data_xy(f):
    result=core.pysorted_data_set();
    
    for i in f:
        x,y=i.split()
        x=float(x)
        y=float(y)
        xe1=0
        xe2=0
        ye1=0
        ye2=0

        result.add_data(core.pydata(x,y,ye1,ye2,xe1,xe2))
    
    return result;


def show_data(f):
    ndata=f.get_data_set().size()
    for i in range(0,ndata):
        d=f.get_data_set().get_data(i)
        print `d.get_x()`+" "+`d.get_y()`


def show(f):
    nparams=f.get_num_params()
    for i in range(0,nparams):
        print f.get_param_info(i).get_name()+"="+`f.get_param_info(i).get_value()`+" "+f.report_param_status(f.get_param_info(i).get_name())


def bootstrap(f,cnt=100):
    bp=pybootstrap()
    bp.set_fitter(f)
    nparams=f.get_num_params()
    origin_param=f.get_all_params()
    result=numpy.zeros([cnt,nparams])
    for i in range(0,cnt):
        p=bp.sample()
        for j in range(0,nparams):
            result[i,j]=p[j]

    for i in range(0,nparams):
        pname=f.get_param_info(i).get_name()
        f.set_param_value(pname,origin_param[i])
    return result


def plot(f):
    import pylab
    ndata=f.get_data_set().size()
    ninsert=10
    x_data=numpy.zeros(ndata)
    x_model=numpy.zeros(ndata*ninsert-ninsert)
    x_error=numpy.zeros((2,ndata))

    y_data=numpy.zeros(ndata)
    y_model=numpy.zeros(ndata*ninsert-ninsert)
    y_error=numpy.zeros((2,ndata))
#    m=f.model()
    for i in range(0,ndata-1):
        x_data[i]=f.get_data_set().get_data(i).get_x()
        y_data[i]=f.get_data_set().get_data(i).get_y()
        y_error[0,i]=f.get_data_set().get_data(i).get_y_lower_err()
        y_error[1,i]=f.get_data_set().get_data(i).get_y_upper_err()
        x_error[0,i]=f.get_data_set().get_data(i).get_x_lower_err()
        x_error[1,i]=f.get_data_set().get_data(i).get_x_upper_err()

#        print x_data[i],y_data[i]

        x1=x_data[i]
        x2=f.get_data_set().get_data(i+1).get_x()
        
        for j in range(0,ninsert):
            x_model[i*ninsert+j]=x1+(x2-x1)/ninsert*j
            y_model[i*ninsert+j]=f.eval_model_raw(x_model[i*ninsert+j],f.get_all_params())

    x_data[ndata-1]=f.get_data_set().get_data(ndata-1).get_x()
    y_data[ndata-1]=f.get_data_set().get_data(ndata-1).get_y()
    x_error[0,ndata-1]=f.get_data_set().get_data(ndata-1).get_x_lower_err()
    x_error[1,ndata-1]=f.get_data_set().get_data(ndata-1).get_x_upper_err()
    y_error[0,ndata-1]=f.get_data_set().get_data(ndata-1).get_y_lower_err()
    y_error[1,ndata-1]=f.get_data_set().get_data(ndata-1).get_y_upper_err()
#    print x_data,y_data
    pylab.errorbar(x_data,y_data,y_error,x_error,'.')
    pylab.plot(x_model,y_model)
    

def plot_to_svg(f,fpath):
    import matplotlib
    matplotlib.use('svg')
    import matplotlib.pyplot as plt
    ndata=f.get_data_set().size()
    ninsert=10
    x_data=numpy.zeros(ndata)
    x_model=numpy.zeros(ndata*ninsert-ninsert)
    x_error=numpy.zeros((2,ndata))

    y_data=numpy.zeros(ndata)
    y_model=numpy.zeros(ndata*ninsert-ninsert)
    y_error=numpy.zeros((2,ndata))
#    m=f.model()
    for i in range(0,ndata-1):
        x_data[i]=f.get_data_set().get_data(i).get_x()
        y_data[i]=f.get_data_set().get_data(i).get_y()
        y_error[0,i]=f.get_data_set().get_data(i).get_y_lower_err()
        y_error[1,i]=f.get_data_set().get_data(i).get_y_upper_err()
        x_error[0,i]=f.get_data_set().get_data(i).get_x_lower_err()
        x_error[1,i]=f.get_data_set().get_data(i).get_x_upper_err()

#        print x_data[i],y_data[i]

        x1=x_data[i]
        x2=f.get_data_set().get_data(i+1).get_x()
        
        for j in range(0,ninsert):
            x_model[i*ninsert+j]=x1+(x2-x1)/ninsert*j
            y_model[i*ninsert+j]=f.eval_model_raw(x_model[i*ninsert+j],f.get_all_params())

    x_data[ndata-1]=f.get_data_set().get_data(ndata-1).get_x()
    y_data[ndata-1]=f.get_data_set().get_data(ndata-1).get_y()
    x_error[0,ndata-1]=f.get_data_set().get_data(ndata-1).get_x_lower_err()
    x_error[1,ndata-1]=f.get_data_set().get_data(ndata-1).get_x_upper_err()
    y_error[0,ndata-1]=f.get_data_set().get_data(ndata-1).get_y_lower_err()
    y_error[1,ndata-1]=f.get_data_set().get_data(ndata-1).get_y_upper_err()
#    print x_data,y_data

    fig=plt.figure()
    ax=fig.add_subplot(111)
    
    ax.errorbar(x_data,y_data,y_error,x_error,'.')
    ax.plot(x_model,y_model)
    fig.savefig(fpath)
    

    
def dump_qdp(f,target,logx=True,logy=True):
    ndata=f.get_data_set().size()
    ninsert=10
    x_data=numpy.zeros(ndata)
    x_model=numpy.zeros(ndata*ninsert-ninsert)
    x_error=numpy.zeros((2,ndata))
    
    y_data=numpy.zeros(ndata)
    y_model=numpy.zeros(ndata*ninsert-ninsert)
    y_error=numpy.zeros((2,ndata))
#    m=f.model()

    target.write("read terr 1 2\n")
    target.write("skip single\n")
    if logx:
        target.write("log x\n")
    if logy:
        target.write("log y\n")

    target.write("line on 2\n")
    
    for i in range(0,ndata-1):
        x_data[i]=f.get_data_set().get_data(i).get_x()
        y_data[i]=f.get_data_set().get_data(i).get_y()
        y_error[0,i]=f.get_data_set().get_data(i).get_y_lower_err()
        y_error[1,i]=f.get_data_set().get_data(i).get_y_upper_err()
        x_error[0,i]=f.get_data_set().get_data(i).get_x_lower_err()
        x_error[1,i]=f.get_data_set().get_data(i).get_x_upper_err()
        
#        print x_data[i],y_data[i]
        
        x1=x_data[i]
        x2=f.get_data_set().get_data(i+1).get_x()
        
        for j in range(0,ninsert):
            x_model[i*ninsert+j]=x1+(x2-x1)/ninsert*j
            y_model[i*ninsert+j]=f.eval_model_raw(x_model[i*ninsert+j],f.get_all_params())
            
    x_data[ndata-1]=f.get_data_set().get_data(ndata-1).get_x()
    y_data[ndata-1]=f.get_data_set().get_data(ndata-1).get_y()
    x_error[0,ndata-1]=f.get_data_set().get_data(ndata-1).get_x_lower_err();
    x_error[1,ndata-1]=f.get_data_set().get_data(ndata-1).get_x_upper_err();
    y_error[0,ndata-1]=f.get_data_set().get_data(ndata-1).get_y_lower_err();
    y_error[1,ndata-1]=f.get_data_set().get_data(ndata-1).get_y_upper_err();

    for i in range(0,ndata):
        target.write("%E -%E %E %E -%E %E\n"%(x_data[i],x_error[0,i],x_error[1,i],y_data[i],y_error[0,i],y_error[1,i]))


    target.write("no no no\n")
    for i in range(0,len(x_model)):
        target.write("%E 0 0 %E 0 0\n"%(x_model[i],y_model[i]))
    
pass
